package  com.xingcheng.discord.midjourney;

import cn.hutool.core.thread.ThreadUtil;
import com.neovisionaries.ws.client.*;
import eu.bitwalker.useragentutils.UserAgent;
import lombok.extern.slf4j.Slf4j;
import net.dv8tion.jda.api.GatewayEncoding;
import net.dv8tion.jda.api.exceptions.ParsingException;
import net.dv8tion.jda.api.requests.CloseCode;
import net.dv8tion.jda.api.utils.Compression;
import net.dv8tion.jda.api.utils.data.DataObject;
import net.dv8tion.jda.api.utils.data.DataType;
import net.dv8tion.jda.internal.requests.WebSocketCode;
import net.dv8tion.jda.internal.utils.IOUtil;
import net.dv8tion.jda.internal.utils.compress.Decompressor;
import net.dv8tion.jda.internal.utils.compress.ZlibDecompressor;
import org.apache.commons.lang3.StringUtils;

import java.io.IOException;
import java.net.Socket;
import java.net.SocketException;
import java.net.SocketTimeoutException;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.zip.DataFormatException;
/**
 * @author xingchen
 * @version 1.0.0
 * @title WebSocketClient
 * @description
 * @create 2023/12/31 17:02
 **/
@Slf4j
public class WebSocketClient extends WebSocketAdapter {

    protected static final String INVALIDATE_REASON = "INVALIDATE_SESSION";





    protected final Compression compression;
    protected final GatewayEncoding encoding;

    public WebSocket socket;
    protected volatile String sessionId = null;
    protected Decompressor decompressor;
    protected String resumeUrl = null;

    protected final ScheduledExecutorService heartExecutor;




    protected long responseTotal;

    private static final int CONNECT_RETRY_LIMIT = 3;


    protected final AtomicInteger messagesSent = new AtomicInteger(0);

    protected boolean connected = false;




    private Future<?> heartbeatTask;

    private String userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/107.0.0.0 Safari/537.36";

    public static final String gatewayWssUrl = "wss://gateway.discord.gg";

    private final EventMessageListener eventListener;
    private final String host;
    private final Integer port;
    /**
     * use  user token
     */
    public  final String token;


    public WebSocketClient(String host,Integer port,String token,EventMessageListener eventListener) {
        this.compression = Compression.ZLIB;
        this.encoding = GatewayEncoding.JSON;
        this.host=host;
        this.port=port;
        this.token=token;
        this.eventListener=eventListener;
        this.heartExecutor = Executors.newSingleThreadScheduledExecutor();
    }


    /**
     * 是否连接
     * @return
     */
    public boolean isConnected() {
        return connected;
    }



    /**
     * 发送消息
     * @param message
     * @return
     */
    protected boolean send(DataObject message) {
            if (!connected) {
                return false;
            }
            log.trace("<- {}", message);
            socket.sendText(message.toString());
            this.messagesSent.getAndIncrement();
            return true;
    }

    /**
     * 创建websocket工厂
     * @param host
     * @param port
     * @return
     */
    public WebSocketFactory createWebSocketFactory(String host,Integer port) {
        WebSocketFactory webSocketFactory = new WebSocketFactory().setConnectionTimeout(10000);
        if (StringUtils.isNotBlank(host)) {
            ProxySettings proxySettings = webSocketFactory.getProxySettings();
            proxySettings.setHost(host);
            proxySettings.setPort(port);
        }
        return webSocketFactory;
    }

    /**
     * 连接
     */
    public synchronized void connect() {
        try {
            String gatewayUrl = resumeUrl != null ? resumeUrl : gatewayWssUrl;
            gatewayUrl = IOUtil.addQuery(gatewayUrl,
                    "encoding", encoding.name().toLowerCase(),
                    "v", 9
            );
            if (compression != Compression.NONE) {
                gatewayUrl = IOUtil.addQuery(gatewayUrl, "compress", compression.getKey());
                switch (compression) {
                    case ZLIB:
                            decompressor = new ZlibDecompressor(2048);
                        break;
                    default:
                        throw new IllegalStateException("Unknown compression");
                }
            }
            WebSocketFactory socketFactory = new WebSocketFactory(createWebSocketFactory(host,port));
            IOUtil.setServerName(socketFactory, gatewayUrl);
            socket = socketFactory.createSocket(gatewayUrl);
            socket.addHeader("Accept-Encoding", "gzip, deflate, br")
                    .addHeader("Accept-Language", "zh-CN,zh;q=0.9")
                    .addHeader("Cache-Control", "no-cache")
                    .addHeader("Pragma", "no-cache")
                    .addHeader("Sec-Websocket-Extensions", "permessage-deflate; client_max_window_bits")
                    .addHeader("User-Agent", userAgent)
                    .addListener(this)
                    .connect();
        } catch (IOException | WebSocketException | IllegalArgumentException e) {
            resumeUrl = null;
            //Completely fail here. We couldn't make the connection.
            throw new IllegalStateException(e);
        }
    }

    /**
     * 连接成功
     * @param websocket
     * @param headers
     */
    @Override
    public void onConnected(WebSocket websocket, Map<String, List<String>> headers) {
        log.info("[gateway] Connected to websocket");
        this.connected = true;
        if (sessionId == null) {
            //发送身份认证
            sendIdentify();
        }else {
            //重连
            sendResume();
        }
    }
    /**
     * 连接关闭
     * @param websocket
     * @param serverCloseFrame
     * @param clientCloseFrame
     * @param closedByServer
     */
    @Override
    public void onDisconnected(WebSocket websocket, WebSocketFrame serverCloseFrame, WebSocketFrame clientCloseFrame, boolean closedByServer) {
        reset();
        Integer code = CloseCode.GRACEFUL_CLOSE.getCode();
        String closeReason = "";
        if (clientCloseFrame != null) {
            code = clientCloseFrame.getCloseCode();
            closeReason = clientCloseFrame.getCloseReason();
        } else if (serverCloseFrame != null) {
            code = serverCloseFrame.getCloseCode();
            closeReason = serverCloseFrame.getCloseReason();
        }
        if (code >= CloseCode.INVALID_SHARD.getCode() || code.equals(CloseCode.AUTHENTICATION_FAILED.getCode())) {
            log.warn("[gateway] Websocket closed and can't reconnect! code: {}, reason: {}", code, closeReason);
            System.exit(code);
            return;
        }
        log.warn("[gateway] Websocket closed and will be reconnect... code: {}, reason: {}", code, closeReason);
        ThreadUtil.execute(() -> {
            try {
                reconnect(0);
            } catch (Exception e) {
                log.error("[gateway] Websocket reconnect error", e);
                System.exit(1);
            }
        });
    }

    /**
     * 重试启动
     * @param currentRetryTime
     * @throws Exception
     */
    private void reconnect(int currentRetryTime) throws Exception {
        try {
            connect();
        } catch (Exception e) {
            if (currentRetryTime < CONNECT_RETRY_LIMIT) {
                currentRetryTime++;
                log.warn("[gateway] Websocket reconnect fail, retry {} time... error: {}", currentRetryTime,
                        e.getMessage());
                Thread.sleep(5000L);
                reconnect(currentRetryTime);
            } else {
                throw e;
            }
        }
    }

    /**
     * 重置
     */
    private void reset() {
        this.connected = false;
        this.sessionId = null;
        this.decompressor = null;
        this.socket = null;
        if (this.heartbeatTask != null) {
            this.heartbeatTask.cancel(true);
            this.heartbeatTask = null;
        }
    }



    /**
     * 发送心跳
     */
   protected void sendKeepAlive() {
        DataObject keepAlivePacket =
                DataObject.empty()
                        .put("op", WebSocketCode.HEARTBEAT)
                        .put("d", responseTotal);
        send(keepAlivePacket);
    }

    /**
     * 设置心跳
     * @param timeout
     */
    protected void setupKeepAlive(int timeout)
    {
        try
        {
            Socket rawSocket = this.socket.getSocket();
            if (rawSocket != null) {
                // setup a timeout when we miss heartbeats
                rawSocket.setSoTimeout(timeout + 10000);
            }
        }
        catch (SocketException ex)
        {
            log.warn("Failed to setup timeout for socket", ex);
        }
        if (this.heartbeatTask == null && this.heartExecutor != null) {
            heartbeatTask = heartExecutor.scheduleAtFixedRate(() ->
            {
                if (connected) {
                    sendKeepAlive();
                }
            }, 0, timeout, TimeUnit.MILLISECONDS);
        }
    }


    /**
     * 获取完整的状态
     * @return
     */
    public DataObject getFullPresence() {
        return DataObject.empty()
                .put("afk", false)
                .put("since", System.currentTimeMillis())
                .put("activities", Collections.emptyList())
                .put("status", "online");
    }

    /**
     * 发送认证
     */
    protected void sendIdentify() {
        log.debug("Sending Identify-packet...");
        UserAgent agent = UserAgent.parseUserAgentString(userAgent);
        DataObject connectionProperties = DataObject.empty().put("os", agent.getOperatingSystem().getName())
                .put("browser", agent.getBrowser().getGroup().getName()).put("device", "").put("system_locale", "zh-CN")
                .put("browser_version", agent.getBrowserVersion().toString()).put("browser_user_agent", getToken());
        DataObject payload = DataObject.empty()
                .put("presence", getFullPresence())
                .put("token", getToken())
                .put("properties", connectionProperties);
        DataObject identify = DataObject.empty()
                .put("op", WebSocketCode.IDENTIFY)
                .put("d", payload);

        send(identify);
    }

    /**
     * 获取token
     * @return
     */
    public String getToken(){
        return token;
    }

    /**
     * 发送恢复
     */
    protected void sendResume() {
        log.debug("Sending Resume-packet...");
        DataObject resume = DataObject.empty()
                .put("op", WebSocketCode.RESUME)
                .put("d", DataObject.empty()
                        .put("session_id", sessionId)
                        .put("token", getToken())
                        .put("seq", responseTotal));
        send(resume);
    }
    protected void handleEvent(DataObject content) {
        try {
            onEvent(content);
        } catch (Exception ex) {
            log.error("Encountered exception on lifecycle level\nJSON: {}", content, ex);
        }
    }

    /**
     * 事件处理
     * @param content
     */
    protected void onEvent(DataObject content) {
        int opCode = content.getInt("op");
        if (!content.isNull("s")) {
            responseTotal=content.getInt("s");
        }
        switch (opCode) {
            case WebSocketCode.DISPATCH:
                onDispatch(content);
                break;
            case WebSocketCode.HEARTBEAT:
                log.debug("Got Keep-Alive request (OP 1). Sending response...");
                sendKeepAlive();
                break;
            case WebSocketCode.RECONNECT:
                log.debug("Got Reconnect request (OP 7). Closing connection now...");
                close(CloseCode.RECONNECT.getCode(), "OP 7: RECONNECT");
                break;
            case WebSocketCode.INVALIDATE_SESSION:
                log.debug("Got Invalidate request (OP 9). Invalidating...");
                close(CloseCode.GRACEFUL_CLOSE.getCode(), INVALIDATE_REASON);
                break;
            case WebSocketCode.HELLO:
                log.debug("Got HELLO packet (OP 10). Initializing keep-alive.");
                final DataObject data = content.getObject("d");
                setupKeepAlive(data.getInt("heartbeat_interval"));
                break;
            case WebSocketCode.HEARTBEAT_ACK:
                log.trace("Got Heartbeat Ack (OP 11).");
                break;
            default:
                log.debug("Got unknown op-code: {} with content: {}", opCode, content);
        }
    }
    /**
     * 事件分发
     * @param raw
     */
    protected void onDispatch(DataObject raw) {
        String type = raw.getString("t");
        if (!raw.isType("d", DataType.OBJECT)) {
            // Needs special handling due to content of "d" being an array
            return;
        }
        DataObject content = raw.getObject("d");
        log.trace("{} -> {}", type, content);
        try {
            if ("READY".equals(type)) {
                //discord 交互时，把sessionId 记录下来,放入请求数据中
                this.sessionId = content.getString("session_id");
                this.resumeUrl = content.getString("resume_gateway_url");
                this.connected=true;
                return;
            }
            try {
                this.eventListener.onMessage(raw);
            } catch (Exception e) {
                log.error("eventListener error", e);
            }
        } catch (ParsingException ex) {
            log.warn("Got an unexpected Json-parse error. Please redirect the following message to the devs:\n\t {}\n\t{}\n\t{} -> {}",
                    9, ex.getMessage(), type, content, ex);
        } catch (Exception ex) {
            log.error("Got an unexpected error. Please redirect the following message to the devs:\n\t {}\n\t{} -> {}",
                    9, type, content, ex);
        }
    }

    /**
     * 文本消息
     * @param websocket
     * @param data
     */
    @Override
    public void onTextMessage(WebSocket websocket, byte[] data) {
        handleEvent(DataObject.fromJson(data));
    }

    /**
     * 二进制消息
     * @param websocket
     * @param binary
     * @throws DataFormatException
     */
    @Override
    public void onBinaryMessage(WebSocket websocket, byte[] binary) throws DataFormatException
    {
        DataObject message = handleBinary(binary);
        if (message != null) {
            handleEvent(message);
        }
    }

    /**
     * 转换字节码
     * @param binary
     * @return
     * @throws DataFormatException
     */
    protected DataObject handleBinary(byte[] binary) throws DataFormatException {
        // Scoping allows us to print the json that possibly failed parsing
        byte[] data;
        try {
            data = decompressor.decompress(binary);
            if(data == null){
                return null;
            }
        } catch (DataFormatException e) {
            log.error("4900", "MALFORMED_PACKAGE_Failed");
            throw e;
        }
        try{
            return DataObject.fromJson(data);
        }catch (Exception e) {
            log.error("Failed to parse json: {}", e.getMessage());
            throw e;
        }
    }

    @Override
    public void onError(WebSocket websocket, WebSocketException cause) throws Exception
    {
        if (cause.getCause() instanceof SocketTimeoutException) {
            log.debug("Socket timed out");
        }
        else if (cause.getCause() instanceof IOException) {
            log.debug("Encountered I/O error", cause);
        } else {
            log.error("There was an error in the WebSocket connection", cause);
        }
    }

    /**
     * 关闭前置操作
     */
    private void prepareClose() {
        try {
            if (socket != null){
                Socket rawSocket = this.socket.getSocket();
                if (rawSocket != null) {
                    rawSocket.setSoTimeout(10000);
                }
            }
        }catch (SocketException ignored) {

        }
    }

    /**
     * 关闭
     */
    public void close() {
        prepareClose();
        if (socket != null) {
            socket.sendClose(1000);
        }
    }

    /**
     * 关闭
     * @param code
     */
    public void close(int code) {
        prepareClose();
        if (socket != null) {
            socket.sendClose(code);
        }
    }

    /**
     * 关闭
     * @param code
     * @param reason
     */
    public void close(int code, String reason) {
        prepareClose();
        if (socket != null) {
            socket.sendClose(code, reason);
        }
    }
}

